#!/usr/bin/env python

import sys, os, re, time
import pcap
import string
import socket
import struct
import urllib2
from DNS.Lib import DnsResult, Munpacker

re_mac = re.compile('^%s$' % ':'.join(['[a-f0-9]{2}' for x in range(6)]), re.I)
re_ip4 = re.compile('^%s$' % '.'.join(['\d{1,3}' for x in range(4)]))
re_printable = re.compile('([ -~\s]{6,})')
re_oui = re.compile('([A-F0-9]{2}-[A-F0-9]{2}-[A-F0-9]{2})\s+\(hex\)\s+(\S+.*)', re.I)
re_drop = re.compile(r',? ?(ltd|inc|plc|llc|limited|co|corp|co\., ltd|corporation)\.?$', re.I)
re_http_req = re.compile(r'([A-Z]{3,}) (\S+) HTTP/\d', re.I)

def dumphex(s, octets_per_line=32):
    bytes = map(lambda x: '%02x' % ord(x), s)
    for i in xrange(0, len(bytes)+1, octets_per_line):
        print '  %s' % string.join(bytes[i:(i+octets_per_line)], ' ')

def ip2name(ip):
    if 'A' not in Sniffer.DnsAnswer: return
    dnsa = Sniffer.DnsAnswer['A']
    if ip not in dnsa: return
    name = dnsa[ip]
    if 'CNAME' not in Sniffer.DnsAnswer: return name
    dnsc = Sniffer.DnsAnswer['CNAME']
    if name not in dnsc: return name
    return dnsc[name]

class Sniffer:
    Dir = os.path.join(os.getenv('HOME'), 'NetTee')
    if not os.path.isdir(Dir):
        os.mkdir(Dir)
    oui_url = 'http://standards.ieee.org/regauth/oui/oui.txt'
    TcpStream = {}
    RequestUrl = {}
    ContentLength = {}
    TcpDir = os.path.join(Dir, 'tcp')
    if not os.path.isdir(TcpDir):
        os.mkdir(TcpDir)
    DnsAnswer = {}
    OuiVendor = {}

    def __init__(self, dev, snaplen=1600, promisc=False, to_ms=100, noself=True):
        self.get_oui_vendors()
        self.dev = dev
        self.get_dev_addr(dev)
        obj = pcap.pcapObject()
        obj.open_live(dev, snaplen, promisc, to_ms)
        if promisc:
            print 'INFO: %s is in promiscuous mode.' % dev
        tcp_filter = 'tcp or (udp and port 53)'
        if noself:
            tcp_filter += ' and not ether host %s' % self.lladdr
        print 'Filter:', tcp_filter
        obj.setfilter(tcp_filter, 0, 0)
        self.promisc = promisc
        self.pcap = obj

    def get_oui_vendors(self):
        oui_file = os.path.join(self.Dir, 'oui.txt')
        if not os.path.exists(oui_file):
            print 'INFO: getting oui.txt from', self.oui_url
            try:
                http = urllib2.urlopen(self.oui_url)
            except:
                'WARN: failed to get oui.txt'
                return
            oui_w = open(oui_file, 'w')
            for line in http.readlines():
                found = re_oui.match(line)
                if found:
                    oui_w.write('%s %s\n' %(found.group(1).lower().replace('-', ':'), re_drop.sub('', found.group(2))))
            oui_w.close()
        try:
            oui_r = open(oui_file, 'r')
        except:
            print 'WARN: failed to open', oui_file, sys.exc_info()[1]
            return
        for line in oui_r.readlines():
            line = line.rstrip()
            Sniffer.OuiVendor[line[:8]] = line[9:]

    def get_dev_addr(self, dev):
        lladdr = ""
        ipaddr = ""
        for ifcfg in pcap.findalldevs():
            if ifcfg[0] != dev: continue
            for addr_tuple in ifcfg[2]:
                if not lladdr and re_mac.match(addr_tuple[0]):
                    lladdr = addr_tuple[0]
                elif not ipaddr and re_ip4.match(addr_tuple[0]):
                    ipaddr = addr_tuple[0]
                if lladdr and ipaddr: break
            break
        if lladdr:
            print "INFO: %s lladdr is %s" % (dev, lladdr)
        else:
            print "ERROR: unable to find lladdr for %s" % dev
        if ipaddr:
            if 'A' not in self.DnsAnswer:
                self.DnsAnswer['A'] = { ipaddr: 'MyIP' }
            else:
                self.DnsAnswer['A'][ipaddr] = 'MyIP'
            print 'INFO: %s IP addr is %s' %(dev, ipaddr)
        else:
            print 'INFO: %s has no IP addr.' % dev
        self.lladdr = lladdr
        self.ipaddr = ipaddr

    def dispatch(self, limit=500):
        obj = self.pcap
        callback = self.sniff
        pkt_count = 0
        try:
            while pkt_count < limit:
                if obj.dispatch(1, callback):
                    pkt_count += 1
        except KeyboardInterrupt:
            if self.promisc:
                print '\nDevice %s was in PROMISC mode.' % self.dev
            print '%s' % sys.exc_type
            print '%d packets received, %d packets dropped, %d packets dropped by interface' % self.pcap.stats()
            stream = self.TcpStream
            print '%d TCP streams in memory' % len(stream.keys())
            for stream_key, frames in stream.items():
                print '%s:%d %s:%d' % stream_key, len(frames), 'segment abandoned'

    def sniff(self, pktlen, data, timestamp):
        frame = Frame(pktlen, data, timestamp)
        if frame.higher:
            l3 = frame.higher
            if l3.higher:
                l4 = l3.higher
                proto = l4.proto
                if proto == 'udp' and (l4.data['src'] == 53 or l4.data['dst'] == 53):
                    #for k, v in Sniffer.DnsAnswer.items(): print k, v
                    return
                if proto == 'tcp':
                    self.update_tcp_stream(frame, l3.data, l4.data)
                    return
                print l4.header(), 'Unknown layer 4'
            else:
                print l3.header(), 'Unknown layer 3'
        else:
            print frame.header(), 'Unknown frame'
        print

    def get_stream_key_and_file(self, ip_data, tcp_data):
        '''Generate stream key for storing frame in memory and create file name for storing frame on file system'''
        src_addr = ip_data['src']
        dst_addr = ip_data['dst']
        src_name = ip2name(src_addr) or src_addr
        dst_name = ip2name(dst_addr) or dst_addr

        src_port = tcp_data['src']
        dst_port = tcp_data['dst']

        tcp_dir = self.TcpDir
        if src_port < dst_port:
            stream_key = (dst_addr, dst_port, src_addr, src_port)
            stream_dir = os.path.join(tcp_dir, dst_addr + '~' + src_addr)
            alias_dir  = os.path.join(tcp_dir, dst_name + '~' + src_name)
            file_name = '%d-%d' %(dst_port, src_port)
        elif src_port > dst_port or src_addr > dst_addr:
            stream_key = (src_addr, src_port, dst_addr, dst_port)
            stream_dir = os.path.join(tcp_dir, src_addr + '~' + dst_addr)
            alias_dir  = os.path.join(tcp_dir, src_name + '~' + dst_name)
            file_name = '%d-%d' %(src_port, dst_port)

        if not os.path.exists(stream_dir): os.mkdir(stream_dir)
        if not os.path.exists(alias_dir):  os.symlink(stream_dir, alias_dir)

        stream_file = os.path.join(stream_dir, file_name)
        return(stream_key, stream_file)

    def get_http_headers(self, printable):
        headers = {}
        for line in printable.split('\r\n'):
            if line.find(': ') > 0:
                f = line.split(': ')
                headers[f[0]] = f[1]
        return headers

    def update_tcp_stream(self, frame, ip_data, tcp_data):
        print frame.higher.higher.header()
        payload_len = tcp_data['payload_len']
        if payload_len == 0:
            # Ignore TCP segments without payload
            print
            return

        (stream_key, stream_file) = self.get_stream_key_and_file(ip_data, tcp_data)
        streams = self.TcpStream

        http_type = self.get_http_type(tcp_data, stream_key)
        if http_type:
            if stream_key in streams and streams[stream_key]:
                # Process all segments in the stream
                self.assemble_http_segments(streams[stream_key], stream_key, stream_file)
                del(streams[stream_key])
                print 'INFO: %d tcp streams in memory\n' % len(streams.keys())
            headers = tcp_data['http_headers']
            header_len = tcp_data['http_header_len']
            clen = payload_len - header_len
            if http_type != 'response' or 'Content-Length' not in headers or clen >= int(headers['Content-Length']):
                # Current frame is a complet HTTP message
                # No need to queue in the stream
                self.assemble_http_segments([frame], stream_key, stream_file)
            else:
                # This is the first segment of a long HTTP message
                #print "DEBUG: first segment of a long HTTP message\n%s\n%s\n" %(frame.higher.higher.header(), tcp_data)
                streams[stream_key] = [frame]
                content_len = int(headers['Content-Length'])
                if stream_key in self.ContentLength:
                    self.ContentLength[stream_key].append(content_len)
                else:
                    self.ContentLength[stream_key] = [content_len]
        elif stream_key in streams:
            streams[stream_key].append(frame)
            print 'DEBUG####: pondering stream_key', stream_key, len(streams[stream_key]), 'segments'
            if stream_key in Sniffer.ContentLength:
                segments = streams[stream_key]
                content_len = Sniffer.ContentLength[stream_key][0]
                total_payload_len = sum([f.higher.higher.data['payload_len'] for f in segments])
                print 'DEBUG####: %d segments, total payload len = %d, content_len = %d' % (len(segments), total_payload_len, content_len)
                if total_payload_len - segments[0].higher.higher.data['http_header_len'] >= content_len:
                    self.assemble_http_segments(segments, stream_key, stream_file)
                    del(streams[stream_key])
                    Sniffer.ContentLength[stream_key].pop()
        else:
            print 'INFO: Non-HTTP packet dropped'
            dumphex(tcp_data['payload'][:128])

    def get_http_type(self, tcp_data, stream_key):
        payload = tcp_data['payload']
        http_type = None
        found = re_printable.match(payload)
        if not found: return
        printable = found.group(1)
        printable_len = len(printable)
        if printable[:5].upper() == 'HTTP/':
            #### HTTP response
            http_type = 'response'
            headers = self.get_http_headers(printable)
            if 'Content-Length' in headers:
                clen = headers['Content-Length']
            else:
                clen = 'unknown'
            print 'INFO: HTTP', http_type, clen, 'bytes'
        else:
            found_req = re_http_req.match(printable)
            if not found_req: return
            ##### HTTP request
            http_type = 'request_' + found_req.group(1)
            request_url = found_req.group(2)
            tcp_data['request_url'] = request_url
            headers = self.get_http_headers(printable)
            host = headers.get('Host', '')
            print 'INFO: HTTP', http_type, host+request_url
            if len(request_url) > 128: request_url = request_url[:125] + '...'
            if stream_key in self.RequestUrl:
                self.RequestUrl[stream_key].append(request_url)
            else:
                self.RequestUrl[stream_key] = [request_url]

        header_len = printable.find('\r\n\r\n')
        if header_len < 0:
            header_len = tcp_data['payload_len']
        else:
            header_len += 4

        tcp_data['http_type'] = http_type
        tcp_data['http_headers'] = headers
        tcp_data['http_header_len'] = header_len
        return http_type

    def assemble_http_segments(self, frames, stream_key, stream_file):
        tcp0 = frames[0].higher.higher
        tcp_data = tcp0.data
        if 'http_type' not in tcp_data:
            print '*** How does this get in here:'
            print tcp0.header(), "\n"
            return
        http_type = tcp_data['http_type']
        http_header_file = '%s.http' %(stream_file)
        fh = open(http_header_file, 'a')
        fh.write(tcp_data['payload'][:tcp_data['http_header_len']]+'\n')
        if http_type != 'response':
            fh.close()
            return

        file = '%s-%x' %(stream_file, int(time.time()))
        try:
            file += Sniffer.RequestUrl[stream_key].pop().replace('/', '~')
        except:
            file += 'UnknownUrl'
        try:
            ext = tcp_data['http_headers']['Content-Type'].split('/')[1]
            if ext == 'jpeg': ext = 'jpg'
        except:
            ext = 'UNK'
        if file[-len(ext):].lower() != ext.lower():
            file += '.' + ext
        try:
            if tcp_data['http_headers']['Content-Encoding'] == 'gzip':
                file += '.gz'
        except:
            pass
        fh.write("\nfile: "+file+"\n\n")
        fh.close()

        fh = open(file, 'w')
        data = tcp_data['payload'][tcp_data['http_header_len']:]
        fh.write(data)
        try:
            content_len = int(tcp_data['http_headers']['Content-Length'])
        except:
            content_len = -1
        byte_count = len(data)
        for frame in frames[1:]:
            #print 'DEBUG: saving packet\n%s' % frame.higher.higher.header()
            payload = frame.higher.higher.data['payload']
            fh.write(payload)
            byte_count += len(payload)
            if content_len >= 0 and byte_count > content_len:
                print 'WARN: Content-Length %d but got %d bytes' %(content_len, byte_count)
        fh.close()
        if content_len >= 0:
            print 'INFO: %d/%d bytes TCP stream saved: %s\n' % (os.path.getsize(file), content_len, file)
        else:
            print 'INFO: %d bytes TCP stream saved: %s\n'  % (os.path.getsize(file), file)

class DNSdata:
    def __init__(self, udp_data):
        data = udp_data.data['payload']
        self.dns = DnsResult(Munpacker(data), None)
        dnsa = Sniffer.DnsAnswer
        for a in self.dns.answers:
            n = a['name']
            t = a['typename']
            d = a['data']
            if t in dnsa:
                dnsa[t][d] = n
            else:
                dnsa[t] = {d: n}

    def dump(self):
        dns = self.dns
        h = dns.header
        print ';; ->>HEADER<<- opcode %s, status %s, id %d'%(h['opcode'],h['status'],h['id'])
        flags=filter(lambda x,h=h:h[x],('qr','aa','rd','ra','tc'))
        print ';; flags: %s; Ques: %d, Ans: %d, Auth: %d, Addit: %d'%(
            string.join(flags), h['qdcount'], h['ancount'], h['nscount'], h['arcount'])
        if dns.questions: print ';; QUESTIONS:',
        for q in dns.questions:
            print '%s, type = %s, class = %s'%(q['qname'], q['qtypestr'], q['qclassstr'])
        if dns.answers: print ';; ANSWERS:',
        for a in dns.answers:
            print '%-20s %-6s  %-6s  %s'%(a['name'], `a['ttl']`, a['typename'], a['data'])
        if dns.additional: print ';; ADDITIONAL RECORDS:',
        for a in dns.additional:
            print '%-20s %-6s  %-6s  %s'%(a['name'], `a['ttl']`, a['typename'], a['data'])
        print

class UDPdata:
    proto = 'udp'
    ProtoClass = { 53: DNSdata }

    def __init__(self, ip_packet):
        self.lower = ip_packet
        data = ip_packet.data['payload']
        src = struct.unpack('>H', data[0:2])[0]
        dst = struct.unpack('>H', data[2:4])[0]
        d = { 'src': src, 'dst': dst }
        d['length'] = struct.unpack('>H', data[4:6])[0]
        d['chksum'] = struct.unpack('>H', data[6:8])[0]
        d['payload'] = data[8:]
        self.data = d
        self.higher = None
        for port in [src, dst]:
            if port in self.ProtoClass and self.ProtoClass[port]:
                self.higher = self.ProtoClass[port](self)
                break

    def header(self):
        d = self.data
        return '%s\nUDP: %d > %d' % (self.lower.header(), d['src'], d['dst'])

class TCPdata:
    proto = 'tcp'
    FlagNames = ['FIN', 'SYN', 'RST', 'PSH', 'ACK', 'URG', 'ECE', 'CWR']
    IgnoredPorts = [22, 443]

    def __init__(self, ip_packet):
        self.lower = ip_packet
        data = ip_packet.data['payload']
        src = struct.unpack('>H', data[0:2])[0]
        dst = struct.unpack('>H', data[2:4])[0]
        data_offset = ord(data[12:13]) >> 4
        if data_offset > 5:
            options = data[20:data_offset*4]
        else:
            options = None
        d = { 'src': src, 'dst': dst }
        d['seq'] = struct.unpack('>L', data[4:8])[0]
        d['ack'] = struct.unpack('>L', data[8:12])[0]
        d['flags'] = self.flags(data[13:14])
        d['window'] = struct.unpack('>H', data[14:16])[0]
        d['option'] = options
        if src in self.IgnoredPorts or dst in self.IgnoredPorts:
            d['payload'] = ''
            d['payload_len'] = 0
        else:
            d['payload'] = data[(data_offset*4):]
            d['payload_len'] = len(d['payload'])
        self.higher = None
        self.data = d

    def flags(self, octet):
        value = ord(octet)
        res = []
        f = 1
        for name in self.FlagNames:
            if f & value: res.append(name)
            f <<= 1
        return res

    def header(self):
        d = self.data
        return '%s\nTCP: %d > %d seq %d ack %d window %d flags %s payload %d' %(self.lower.header(),
            d['src'], d['dst'], d['seq'], d['ack'], d['window'], ' '.join(d['flags']), d['payload_len'])

class IPdata:
    ProtoClass = {
        socket.IPPROTO_TCP:  TCPdata,
        socket.IPPROTO_UDP:  UDPdata,
        socket.IPPROTO_ICMP: None
        }

    def __init__(self, frame):
        self.lower = frame
        s = frame.data['payload']
        d = {}
        d['version'] = (ord(s[0]) & 0xf0) >> 4
        d['header_len'] = ord(s[0]) & 0x0f
        d['tos'] = ord(s[1])
        d['total_len'] = socket.ntohs(struct.unpack('H',s[2:4])[0])
        d['id'] = socket.ntohs(struct.unpack('H',s[4:6])[0])
        d['flags'] = (ord(s[6]) & 0xe0) >> 5
        d['fragment_offset'] = socket.ntohs(struct.unpack('H',s[6:8])[0] & 0x1f)
        d['ttl'] = ord(s[8])
        d['proto'] = ord(s[9])
        d['checksum'] = socket.ntohs(struct.unpack('H',s[10:12])[0])
        d['src'] = pcap.ntoa(struct.unpack('i',s[12:16])[0])
        d['dst'] = pcap.ntoa(struct.unpack('i',s[16:20])[0])
        if d['header_len'] > 5:
            d['options'] = s[20:4*(d['header_len']-5)]
        else:
            d['options'] = None
        d['payload'] = s[4*d['header_len']:]
        self.data = d
        self.higher = None
        proto = d['proto']
        if proto in self.ProtoClass and self.ProtoClass[proto]:
            self.higher = self.ProtoClass[proto](self)

    def header(self):
        d = self.data
        src = d['src']
        src_name = ip2name(src)
        if src_name:
            src += ' (%s)' % src_name
        dst = d['dst']
        dst_name = ip2name(dst)
        if dst_name:
            dst += ' (%s)' % dst_name
        if d['flags'] != 2:
            flags = 'flags: %d' % d['flags']
        else:
            flags = ''
        return '%s\nIP %d: %s > %s size: %d %s' %(self.lower.header(), d['version'], src, dst, d['total_len'], flags)

class Frame:
    TypeClass = {'0800': IPdata}

    def __init__(self, size, data, timestamp):
        dst_mac = self.pp_mac(data[0:6])
        src_mac = self.pp_mac(data[6:12])
        type = '%04x' % struct.unpack('>H', data[12:14])
        d = { 'size': size, 'timestamp': timestamp, 'src': src_mac, 'dst': dst_mac }
        d['type'] = type
        d['payload'] = data[14:]
        self.data = d
        self.higher = None
        if type in self.TypeClass and self.TypeClass[type]:
            self.higher = self.TypeClass[type](self)

    def pp_mac(self, s):
        mac = ':'.join(['%02x' % ord(x) for x in s])
        if mac[:8] in Sniffer.OuiVendor:
            mac += ' (%s)' % Sniffer.OuiVendor[mac[:8]]
        return mac

    def header(self):
        d = self.data
        src = d['src']
        dst = d['dst']
        timestamp = d['timestamp']
        hms = time.strftime('%F %H:%M:', time.localtime(timestamp)) + ('%09.6f' % (timestamp % 60))
        return '%s MAC: %s > %s size: %d' %(hms, src, dst, d['size'])

if __name__=='__main__':
    if len(sys.argv) == 1:
        print 'Usage: sniff.py <interface>'
        sys.exit(0)
    dev = sys.argv[1]
    sniffer = Sniffer(dev, promisc=True, noself=False)
    sniffer.dispatch(100000)
